from flask import Flask, jsonify, request, send_file
from flask_cors import CORS
from flask_jwt_extended import JWTManager, create_access_token, jwt_required, get_jwt_identity
from models import db, User, ExtractionRecord, KnowledgeEntity, KnowledgeRelation, Conversation, Message, DictionaryTerm, Document
from config import Config
from rag_service import get_rag_service
from document_parser import parse_document, get_file_type, is_supported_file
from werkzeug.utils import secure_filename
import json
import os

app = Flask(__name__)
app.config.from_object(Config)

# 配置文件上传
UPLOAD_FOLDER = os.path.join(os.path.dirname(__file__), 'uploads')
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB 最大文件大小

# 初始化扩展
CORS(app)
db.init_app(app)
jwt = JWTManager(app)

# 知识抽取服务（全局单例）
rag_service = None


def init_rag_service():
    """初始化知识抽取服务"""
    global rag_service
    if rag_service is None:
        rag_service = get_rag_service()


# ==================== 用户认证API ====================

@app.route('/api/auth/register', methods=['POST'])
def register():
    """用户注册"""
    data = request.get_json()
    
    username = data.get('username')
    password = data.get('password')
    
    if not all([username, password]):
        return jsonify({'success': False, 'message': '请提供完整的注册信息'}), 400
    
    # 检查用户是否已存在
    if User.query.filter_by(username=username).first():
        return jsonify({'success': False, 'message': '用户名已存在'}), 400
    
    # 创建新用户
    user = User(username=username)
    user.set_password(password)
    
    db.session.add(user)
    db.session.commit()
    
    return jsonify({
        'success': True,
        'message': '注册成功',
        'user': user.to_dict()
    }), 201


@app.route('/api/auth/login', methods=['POST'])
def login():
    """用户登录"""
    data = request.get_json()
    
    username = data.get('username')
    password = data.get('password')
    
    if not all([username, password]):
        return jsonify({'success': False, 'message': '请提供用户名和密码'}), 400
    
    user = User.query.filter_by(username=username).first()
    
    if not user or not user.check_password(password):
        return jsonify({'success': False, 'message': '用户名或密码错误'}), 401
    
    # 生成访问令牌
    access_token = create_access_token(identity=str(user.id))
    
    return jsonify({
        'success': True,
        'message': '登录成功',
        'token': access_token,
        'user': user.to_dict()
    })


@app.route('/api/auth/me', methods=['GET'])
@jwt_required()
def get_current_user():
    """获取当前用户信息"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    return jsonify({
        'success': True,
        'user': user.to_dict()
    })


# ==================== 用户管理API ====================

@app.route('/api/users', methods=['GET'])
@jwt_required()
def get_users():
    """获取用户列表（仅管理员）"""
    current_user_id = int(get_jwt_identity())
    current_user = User.query.get(current_user_id)
    
    if not current_user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not current_user.is_admin:
        return jsonify({'success': False, 'message': '权限不足'}), 403
    
    users = User.query.all()
    return jsonify({
        'success': True,
        'users': [user.to_dict() for user in users]
    })


@app.route('/api/users/<int:user_id>', methods=['PUT'])
@jwt_required()
def update_user(user_id):
    """更新用户信息（仅管理员）"""
    current_user_id = int(get_jwt_identity())
    current_user = User.query.get(current_user_id)
    
    if not current_user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not current_user.is_admin:
        return jsonify({'success': False, 'message': '权限不足'}), 403
    
    user = User.query.get(user_id)
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    data = request.get_json()
    
    if 'username' in data:
        user.username = data['username']
    if 'is_admin' in data:
        user.is_admin = data['is_admin']
    
    db.session.commit()
    
    return jsonify({
        'success': True,
        'message': '更新成功',
        'user': user.to_dict()
    })


@app.route('/api/users/<int:user_id>', methods=['DELETE'])
@jwt_required()
def delete_user(user_id):
    """删除用户（仅管理员）"""
    current_user_id = int(get_jwt_identity())
    current_user = User.query.get(current_user_id)
    
    if not current_user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not current_user.is_admin:
        return jsonify({'success': False, 'message': '权限不足'}), 403
    
    user = User.query.get(user_id)
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    db.session.delete(user)
    db.session.commit()
    
    return jsonify({
        'success': True,
        'message': '删除成功'
    })


# ==================== 知识抽取API ====================

@app.route('/api/extract', methods=['POST'])
@jwt_required()
def extract_knowledge():
    """知识抽取 - 使用用户自定义提示词"""
    user_id = int(get_jwt_identity())
    data = request.get_json()
    
    prompt = data.get('prompt', '')
    document_id = data.get('document_id')
    
    if not prompt:
        return jsonify({'success': False, 'message': '请输入抽取提示词'}), 400
    
    if not document_id:
        return jsonify({'success': False, 'message': '请选择要抽取的文档'}), 400
    
    # 从数据库获取文档
    document = Document.query.get(document_id)
    if not document:
        return jsonify({'success': False, 'message': '文档不存在'}), 404
    
    # 获取文档内容
    document_content = document.content
    if not document_content:
        return jsonify({'success': False, 'message': '文档内容为空'}), 400
    
    # 初始化RAG服务
    init_rag_service()
    
    # 1. 获取所有领域词典（包含向量）
    dictionary_terms = DictionaryTerm.query.all()
    dict_data = [{'id': t.id, 'term': t.term, 'category': t.category, 'definition': t.definition, 'embedding': t.embedding} 
                 for t in dictionary_terms]
    
    # 2. 准备文档数据（使用文档的embedding）
    doc_chunks = []
    if document.embedding:
        doc_chunks.append({
            'text': document_content,
            'embedding': document.embedding
        })
    else:
        # 如果文档没有embedding，临时生成一个
        embedding = rag_service.generate_embedding(document_content[:2000])
        if embedding:
            doc_chunks.append({
                'text': document_content,
                'embedding': json.dumps(embedding)
            })
    
    # 3. 使用向量检索获取相关上下文
    retrieved_context = rag_service.retrieve_relevant_context(
        query=prompt,
        dictionary_terms=dict_data,
        document_chunks=doc_chunks,
        top_k=5
    )
    
    # 4. 使用RAG服务进行知识抽取（传入检索到的上下文）
    result = rag_service.extract_entities(prompt, document_content, retrieved_context)
    
    if not result['success']:
        return jsonify(result), 500
    
    # ===== 过滤孤立实体 =====
    # 1. 收集所有在关系中出现的实体
    relations_data = result['data'].get('relations', [])
    entities_in_relations = set()
    for relation_data in relations_data:
        entities_in_relations.add(relation_data.get('subject', ''))
        entities_in_relations.add(relation_data.get('object', ''))
    
    # 2. 过滤实体：只保留在关系中出现的实体
    entities_data = result['data'].get('entities', {})
    filtered_entities = {}
    for entity_type, entity_values in entities_data.items():
        filtered_values = [value for value in entity_values if value in entities_in_relations]
        # 只保留非空的实体类型
        if filtered_values:
            filtered_entities[entity_type] = filtered_values
    
    # 更新过滤后的实体数据
    result['data']['entities'] = filtered_entities
    # ===== 过滤完成 =====
    
    # 保存抽取记录（保存文档标题而非完整内容）
    extraction = ExtractionRecord(
        user_id=user_id,
        prompt=prompt,
        document=f"文档：{document.title} (ID:{document.id})",
        result=json.dumps(result['data'], ensure_ascii=False)
    )
    db.session.add(extraction)
    db.session.commit()
    
    # 保存实体（只保存过滤后的实体）
    for entity_type, entity_values in filtered_entities.items():
        for entity_value in entity_values:
            entity = KnowledgeEntity(
                extraction_id=extraction.id,
                entity_type=entity_type,
                entity_value=entity_value
            )
            db.session.add(entity)
    
    # 保存关系
    for relation_data in relations_data:
        relation = KnowledgeRelation(
            extraction_id=extraction.id,
            subject=relation_data.get('subject', ''),
            predicate=relation_data.get('predicate', ''),
            object=relation_data.get('object', '')
        )
        db.session.add(relation)
    
    db.session.commit()
    
    return jsonify({
        'success': True,
        'data': result['data'],
        'extraction_id': extraction.id
    })


@app.route('/api/extract/optimize-prompt', methods=['POST'])
@jwt_required()
def optimize_prompt():
    """优化用户提示词"""
    data = request.get_json()
    
    user_prompt = data.get('prompt', '')
    
    if not user_prompt:
        return jsonify({'success': False, 'message': '请输入要优化的提示词'}), 400
    
    # 获取所有领域词典内容作为参考
    dictionary_terms = DictionaryTerm.query.all()
    dictionary_knowledge = "\n".join([term.to_knowledge_text() for term in dictionary_terms])
    
    # 使用RAG服务优化提示词
    result = rag_service.optimize_prompt(user_prompt, dictionary_knowledge)
    
    return jsonify(result)


# ==================== 数据展示API ====================

@app.route('/api/extractions', methods=['GET'])
@jwt_required()
def get_extractions():
    """获取抽取记录列表"""
    try:
        user_id = int(get_jwt_identity())
        print(f"用户ID: {user_id}")
        
        user = User.query.get(user_id)
        print(f"用户对象: {user}")
        
        if not user:
            return jsonify({'success': False, 'message': '用户不存在'}), 404
        
        # 管理员可以看到所有记录，普通用户只能看到自己的
        if user.is_admin:
            extractions = ExtractionRecord.query.order_by(ExtractionRecord.created_at.desc()).all()
        else:
            extractions = ExtractionRecord.query.filter_by(user_id=user_id).order_by(ExtractionRecord.created_at.desc()).all()
        
        print(f"查询到 {len(extractions)} 条记录")
        
        # 尝试序列化每条记录
        extractions_list = []
        for idx, extraction in enumerate(extractions):
            try:
                extractions_list.append(extraction.to_dict())
            except Exception as e:
                print(f"序列化第 {idx} 条记录时出错: {e}")
                print(f"记录ID: {extraction.id}, user_id: {extraction.user_id}")
                raise
        
        return jsonify({
            'success': True,
            'extractions': extractions_list
        })
    except Exception as e:
        print(f"获取抽取记录列表时出错: {e}")
        import traceback
        traceback.print_exc()
        return jsonify({'success': False, 'message': str(e)}), 500


@app.route('/api/extractions/<int:extraction_id>', methods=['GET'])
@jwt_required()
def get_extraction(extraction_id):
    """获取单个抽取记录详情"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    extraction = ExtractionRecord.query.get(extraction_id)
    if not extraction:
        return jsonify({'success': False, 'message': '记录不存在'}), 404
    
    # 权限检查
    if not user.is_admin and extraction.user_id != user_id:
        return jsonify({'success': False, 'message': '权限不足'}), 403
    
    # 获取关联的实体和关系
    entities = KnowledgeEntity.query.filter_by(extraction_id=extraction_id).all()
    relations = KnowledgeRelation.query.filter_by(extraction_id=extraction_id).all()
    
    return jsonify({
        'success': True,
        'extraction': extraction.to_dict(),
        'entities': [entity.to_dict() for entity in entities],
        'relations': [relation.to_dict() for relation in relations]
    })


@app.route('/api/entities', methods=['GET'])
@jwt_required()
def get_entities():
    """获取所有实体"""
    entities = KnowledgeEntity.query.all()
    return jsonify({
        'success': True,
        'entities': [entity.to_dict() for entity in entities]
    })


@app.route('/api/relations', methods=['GET'])
@jwt_required()
def get_relations():
    """获取所有关系"""
    relations = KnowledgeRelation.query.all()
    return jsonify({
        'success': True,
        'relations': [relation.to_dict() for relation in relations]
    })


# ==================== 知识图谱API ====================

@app.route('/api/knowledge-graph', methods=['GET'])
@jwt_required()
def get_knowledge_graph():
    """获取知识图谱数据 - 只显示有关系连接的实体"""
    # 获取所有实体和关系
    entities = KnowledgeEntity.query.all()
    relations = KnowledgeRelation.query.all()
    
    # ===== 收集所有在关系中出现的实体 =====
    entities_in_relations = set()
    for relation in relations:
        entities_in_relations.add(relation.subject)
        entities_in_relations.add(relation.object)
    # ===== 收集完成 =====
    
    # 构建实体类型映射（只包含有关系的实体）
    entity_type_map = {}
    for entity in entities:
        if entity.entity_value in entities_in_relations:
            entity_type_map[entity.entity_value] = entity.entity_type
    
    # 构建图谱节点（去重，只包含有关系的实体）
    nodes = {}
    
    # 首先添加所有已知实体（过滤孤立实体）
    for entity in entities:
        if entity.entity_value in entities_in_relations and entity.entity_value not in nodes:
            nodes[entity.entity_value] = {
                'id': entity.entity_value,
                'label': entity.entity_value,
                'type': entity.entity_type,
                'category': entity.entity_type
            }
    
    # 然后确保关系中的所有实体也作为节点（推断类型）
    for relation in relations:
        # 添加主体节点
        if relation.subject not in nodes:
            # 根据关系推断类型
            inferred_type = infer_entity_type(relation.predicate, 'subject')
            nodes[relation.subject] = {
                'id': relation.subject,
                'label': relation.subject,
                'type': inferred_type,
                'category': inferred_type
            }
        
        # 添加客体节点
        if relation.object not in nodes:
            # 根据关系推断类型
            inferred_type = infer_entity_type(relation.predicate, 'object')
            nodes[relation.object] = {
                'id': relation.object,
                'label': relation.object,
                'type': inferred_type,
                'category': inferred_type
            }
    
    # 构建图谱边 - 确保方向性：subject -> predicate -> object
    edges = []
    for relation in relations:
        edges.append({
            'source': relation.subject,  # 关系的主体
            'target': relation.object,   # 关系的客体
            'label': relation.predicate,
            'relation': relation.predicate
        })
    
    return jsonify({
        'success': True,
        'graph': {
            'nodes': list(nodes.values()),
            'edges': edges
        }
    })


def infer_entity_type(predicate, position):
    """根据关系类型推断实体类型"""
    # 四种标准关系的类型映射
    relation_type_map = {
        '发生故障': {
            'subject': '设备名',
            'object': '故障状态'
        },
        '表现': {
            'subject': '故障状态',
            'object': '性能表征'
        },
        '由原因引起': {
            'subject': '故障状态',
            'object': '故障原因'
        },
        '通过办法解决': {
            'subject': '故障状态',
            'object': '解决办法'
        }
    }
    
    if predicate in relation_type_map:
        return relation_type_map[predicate].get(position, '未分类')
    
    return '未分类'


# ==================== 知识库管理API ====================

@app.route('/api/knowledge-base/documents', methods=['GET'])
@jwt_required()
def get_knowledge_documents():
    """获取知识库文档列表"""
    documents = rag_service.documents
    
    # 返回文档列表，每个文档附带索引和摘要
    doc_list = []
    for idx, doc in enumerate(documents):
        doc_list.append({
            'id': idx,
            'title': f'文档 {idx + 1}',
            'summary': doc[:100] + '...' if len(doc) > 100 else doc,
            'content': doc
        })
    
    return jsonify({
        'success': True,
        'documents': doc_list
    })


@app.route('/api/knowledge-base/add', methods=['POST'])
@jwt_required()
def add_knowledge_document():
    """添加文档到知识库"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not user.is_admin:
        return jsonify({'success': False, 'message': '权限不足'}), 403
    
    data = request.get_json()
    document = data.get('document', '')
    
    if not document:
        return jsonify({'success': False, 'message': '请提供文档内容'}), 400
    
    rag_service.add_document(document)
    
    return jsonify({
        'success': True,
        'message': '文档已添加到知识库'
    })


@app.route('/api/knowledge-base/search', methods=['POST'])
@jwt_required()
def search_knowledge():
    """搜索知识库"""
    data = request.get_json()
    query = data.get('query', '')
    top_k = data.get('top_k', 5)
    
    if not query:
        return jsonify({'success': False, 'message': '请提供查询内容'}), 400
    
    results = rag_service.retrieve(query, top_k=top_k)
    
    return jsonify({
        'success': True,
        'results': results
    })


# ==================== 文档管理API ====================

@app.route('/api/documents', methods=['GET'])
@jwt_required()
def get_documents():
    """获取文档列表"""
    documents = Document.query.order_by(Document.created_at.desc()).all()
    return jsonify({
        'success': True,
        'documents': [doc.to_dict() for doc in documents]
    })


@app.route('/api/documents/<int:doc_id>', methods=['GET'])
@jwt_required()
def get_document(doc_id):
    """获取文档详情"""
    document = Document.query.get(doc_id)
    if not document:
        return jsonify({'success': False, 'message': '文档不存在'}), 404
    
    return jsonify({
        'success': True,
        'document': document.to_dict(),
        'content': document.content
    })


@app.route('/api/documents/upload', methods=['POST'])
@jwt_required()
def upload_document():
    """上传文档（仅管理员）"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not user.is_admin:
        return jsonify({'success': False, 'message': '权限不足，仅管理员可上传文档'}), 403
    
    if 'file' not in request.files:
        return jsonify({'success': False, 'message': '未提供文件'}), 400
    
    file = request.files['file']
    title = request.form.get('title', '')
    
    if file.filename == '':
        return jsonify({'success': False, 'message': '未选择文件'}), 400
    
    if not is_supported_file(file.filename):
        return jsonify({'success': False, 'message': '不支持的文件格式'}), 400
    
    file_path = None
    try:
        # 保存文件 - 处理中文文件名
        original_filename = file.filename
        
        # 获取文件扩展名
        file_extension = os.path.splitext(original_filename)[1].lower()
        file_type = file_extension.replace('.', '')
        
        if not file_type:
            return jsonify({'success': False, 'message': '无法识别文件类型，请确保文件有正确的扩展名'}), 400
        
        # 生成唯一文件名，保留扩展名
        import uuid
        unique_filename = f"{uuid.uuid4().hex}{file_extension}"
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
        
        # 保存原始文件名到标题（如果未提供标题）
        if not title:
            # 移除扩展名作为默认标题
            title = os.path.splitext(original_filename)[0]
        
        file.save(file_path)
        file_size = os.path.getsize(file_path)
        
        # 解析文档内容
        content = parse_document(file_path, file_type)
        
        if not content:
            os.remove(file_path)
            return jsonify({'success': False, 'message': '文档内容为空或解析失败'}), 500
        
        # 生成文档向量嵌入
        init_rag_service()
        embedding_text = f"{title}\n{content[:2000]}"  # 使用标题和前2000字符生成embedding
        embedding = rag_service.generate_embedding(embedding_text)
        
        # 保存到数据库
        document = Document(
            title=title,
            filename=original_filename,  # 保存原始文件名（包含中文）
            file_type=file_type,
            file_path=file_path,
            content=content,
            file_size=file_size,
            embedding=json.dumps(embedding) if embedding else None,
            upload_user_id=user_id
        )
        
        db.session.add(document)
        db.session.commit()
        
        return jsonify({
            'success': True,
            'message': '文档上传成功',
            'document': document.to_dict()
        }), 201
        
    except Exception as e:
        # 删除已上传的文件
        if file_path and os.path.exists(file_path):
            os.remove(file_path)
        
        error_message = str(e)
        print(f"文档上传错误: {error_message}")
        return jsonify({'success': False, 'message': f'上传失败: {error_message}'}), 500


@app.route('/api/documents/<int:doc_id>', methods=['DELETE'])
@jwt_required()
def delete_document(doc_id):
    """删除文档及其物理文件（仅管理员）"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not user.is_admin:
        return jsonify({'success': False, 'message': '权限不足，仅管理员可删除文档'}), 403
    
    document = Document.query.get(doc_id)
    if not document:
        return jsonify({'success': False, 'message': '文档不存在'}), 404
    
    file_path = document.file_path
    filename = document.filename
    
    try:
        # 先删除物理文件
        if file_path and os.path.exists(file_path):
            try:
                os.remove(file_path)
                print(f"已删除物理文件: {file_path}")
            except Exception as file_error:
                print(f"删除物理文件失败: {file_path}, 错误: {file_error}")
                # 即使文件删除失败，也继续删除数据库记录
        else:
            print(f"物理文件不存在或路径为空: {file_path}")
        
        # 删除数据库记录
        db.session.delete(document)
        db.session.commit()
        
        print(f"文档删除成功: ID={doc_id}, 文件名={filename}")
        
        return jsonify({
            'success': True,
            'message': '文档及物理文件已删除'
        })
    except Exception as e:
        db.session.rollback()
        error_msg = f'删除失败: {str(e)}'
        print(f"文档删除错误: {error_msg}")
        return jsonify({'success': False, 'message': error_msg}), 500


# ==================== 领域词典API ====================

@app.route('/api/dictionary', methods=['GET'])
@jwt_required()
def get_dictionary_terms():
    """获取所有词典词条"""
    terms = DictionaryTerm.query.order_by(DictionaryTerm.term).all()
    return jsonify({
        'success': True,
        'terms': [term.to_dict() for term in terms]
    })


@app.route('/api/dictionary/<int:term_id>', methods=['GET'])
@jwt_required()
def get_dictionary_term(term_id):
    """获取单个词条详情"""
    term = DictionaryTerm.query.get(term_id)
    if not term:
        return jsonify({'success': False, 'message': '词条不存在'}), 404
    
    return jsonify({
        'success': True,
        'term': term.to_dict()
    })


@app.route('/api/dictionary', methods=['POST'])
@jwt_required()
def add_dictionary_term():
    """添加词典词条（仅管理员）"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not user.is_admin:
        return jsonify({'success': False, 'message': '权限不足，仅管理员可添加词条'}), 403
    
    data = request.get_json()
    
    term_text = data.get('term', '').strip()
    category = data.get('category', '').strip()
    definition = data.get('definition', '').strip()
    
    if not all([term_text, category, definition]):
        return jsonify({'success': False, 'message': '请提供完整的词条信息（名称、类别、释义）'}), 400
    
    # 检查词条是否已存在
    existing = DictionaryTerm.query.filter_by(term=term_text, category=category).first()
    if existing:
        return jsonify({'success': False, 'message': '该词条已存在'}), 400
    
    # 创建新词条
    term = DictionaryTerm(
        term=term_text,
        category=category,
        definition=definition,
        examples=None  # 不再使用示例字段
    )
    
    # 生成向量嵌入
    init_rag_service()
    embedding_text = f"{category} {term_text}: {definition}"
    embedding = rag_service.generate_embedding(embedding_text)
    if embedding:
        term.embedding = json.dumps(embedding)
    
    db.session.add(term)
    db.session.commit()
    
    return jsonify({
        'success': True,
        'message': '词条已添加',
        'term': term.to_dict()
    }), 201


@app.route('/api/dictionary/<int:term_id>', methods=['PUT'])
@jwt_required()
def update_dictionary_term(term_id):
    """更新词典词条（仅管理员）"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not user.is_admin:
        return jsonify({'success': False, 'message': '权限不足，仅管理员可编辑词条'}), 403
    
    term = DictionaryTerm.query.get(term_id)
    if not term:
        return jsonify({'success': False, 'message': '词条不存在'}), 404
    
    data = request.get_json()
    
    # 更新字段
    updated = False
    if 'term' in data:
        term.term = data['term'].strip()
        updated = True
    if 'category' in data:
        term.category = data['category'].strip()
        updated = True
    if 'definition' in data:
        term.definition = data['definition'].strip()
        updated = True
    
    # 如果有更新，重新生成向量嵌入
    if updated:
        init_rag_service()
        embedding_text = f"{term.category} {term.term}: {term.definition}"
        embedding = rag_service.generate_embedding(embedding_text)
        if embedding:
            term.embedding = json.dumps(embedding)
    
    db.session.commit()
    
    return jsonify({
        'success': True,
        'message': '词条已更新',
        'term': term.to_dict()
    })


@app.route('/api/dictionary/<int:term_id>', methods=['DELETE'])
@jwt_required()
def delete_dictionary_term(term_id):
    """删除词典词条（仅管理员）"""
    user_id = int(get_jwt_identity())
    user = User.query.get(user_id)
    
    if not user:
        return jsonify({'success': False, 'message': '用户不存在'}), 404
    
    if not user.is_admin:
        return jsonify({'success': False, 'message': '权限不足，仅管理员可删除词条'}), 403
    
    term = DictionaryTerm.query.get(term_id)
    if not term:
        return jsonify({'success': False, 'message': '词条不存在'}), 404
    
    db.session.delete(term)
    db.session.commit()
    
    return jsonify({
        'success': True,
        'message': '词条已删除'
    })


# ==================== 对话API ====================
# 已禁用AI对话功能
# 如需启用，请恢复 Conversation 和 Message 相关的API端点


# ==================== 初始化数据库 ====================

def init_db():
    """初始化数据库"""
    with app.app_context():
        db.create_all()
        
        # 创建默认管理员账户
        admin = User.query.filter_by(username='admin').first()
        if not admin:
            admin = User(
                username='admin',
                is_admin=True
            )
            admin.set_password('admin123')
            db.session.add(admin)
            db.session.commit()
            print("已创建默认管理员账户: admin / admin123")
        
        # 初始化示例词典数据
        if DictionaryTerm.query.count() == 0:
            sample_terms = [
                {
                    'term': '锅炉',
                    'category': '设备名称',
                    'definition': '火电厂的核心设备，通过燃烧燃料产生高温高压蒸汽，为汽轮机提供动力。',
                    'examples': '例如：循环流化床锅炉、煤粉锅炉'
                },
                {
                    'term': '汽轮机',
                    'category': '设备名称',
                    'definition': '将蒸汽的热能转换为机械能的旋转式动力机械设备。',
                    'examples': '例如：凝汽式汽轮机、背压式汽轮机'
                },
                {
                    'term': '爆管',
                    'category': '故障类型',
                    'definition': '锅炉受热面管道因过热、腐蚀、磨损等原因发生破裂的现象。',
                    'examples': '例如：水冷壁爆管、过热器爆管'
                },
                {
                    'term': '振动超标',
                    'category': '故障类型',
                    'definition': '设备运行时振动幅值超过允许范围，可能导致设备损坏。',
                    'examples': '例如：汽轮机轴振超标、风机振动超标'
                },
                {
                    'term': '热效率',
                    'category': '性能参数',
                    'definition': '火电机组将燃料热能转换为电能的效率指标。',
                    'examples': '例如：机组热效率下降可能由漏风、结焦等原因引起'
                },
                {
                    'term': '真空度',
                    'category': '性能参数',
                    'definition': '凝汽器内的真空程度，影响汽轮机的经济性。',
                    'examples': '例如：真空度下降会导致汽轮机功率降低'
                }
            ]
            
            for term_data in sample_terms:
                term = DictionaryTerm(**term_data)
                db.session.add(term)
            
            db.session.commit()
            print(f"已初始化 {len(sample_terms)} 个示例词典词条")


@app.route('/')
def index():
    """首页"""
    return jsonify({
        'message': '火电知识图谱构建平台 API',
        'description': '基于大语言模型和提示学习的火电知识图谱构建系统',
        'version': '1.0.0'
    })


if __name__ == '__main__':
    init_db()
    init_rag_service()  # 启动时初始化RAG服务
    app.run(debug=True, host='0.0.0.0', port=5000)

